import os
from PIL import Image
root_dir = 'week1/PyTorch学习/dataset/train'
target_ants_dirs = 'ants_image'
target_bees_dirs = 'bees_image'
label_dirs = [target_ants_dirs.split('_')[0], target_bees_dirs.split('_')[0]]
output_dirs = ['ants_label', 'bees_label']

img_path_ants =sorted(os.listdir(os.path.join(root_dir, label_dirs[0])))
img_path_bees =sorted(os.listdir(os.path.join(root_dir, label_dirs[1])))

for i in img_path_ants:
    file_name = i.split('.')[0]
    with open(os.path.join(root_dir, output_dirs[0], file_name + '.txt'), 'w') as f:
      f.write(label_dirs[0])

for i in img_path_bees:
    file_name = i.split('.')[0]
    with open(os.path.join(root_dir, output_dirs[1], file_name + '.txt'), 'w') as f:
      f.write(label_dirs[0])